package wvtools;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileFilter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.Reader;
import java.io.StringReader;
import java.io.Writer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;

import org.dom4j.Document;
import org.dom4j.io.SAXReader;

import preprocess.XMLDoc;

import wvtools.extend.WVTXMLDocumentInfo;
import wvtools.extend.WordVectorWriterExtend;
import wvtools.extend.XMLDocumentLoader;
import edu.udo.cs.wvtool.config.WVTConfiguration;
import edu.udo.cs.wvtool.config.WVTConfigurationFact;
import edu.udo.cs.wvtool.generic.charmapper.WVTCharConverter;
import edu.udo.cs.wvtool.generic.inputfilter.TextInputFilter;
import edu.udo.cs.wvtool.generic.inputfilter.XMLInputFilter;
import edu.udo.cs.wvtool.generic.stemmer.PorterStemmerWrapper;
import edu.udo.cs.wvtool.generic.stemmer.WVTStemmer;
import edu.udo.cs.wvtool.generic.tokenizer.WVTTokenizer;
import edu.udo.cs.wvtool.generic.vectorcreation.TFIDF;
import edu.udo.cs.wvtool.generic.vectorcreation.WVTVectorCreator;
import edu.udo.cs.wvtool.generic.wordfilter.WVTWordFilter;
import edu.udo.cs.wvtool.main.WVTDocumentInfo;
import edu.udo.cs.wvtool.main.WVTFileInputList;
import edu.udo.cs.wvtool.main.WVTWordVector;
import edu.udo.cs.wvtool.main.WVTool;
import edu.udo.cs.wvtool.util.TokenEnumeration;
import edu.udo.cs.wvtool.util.WVToolException;
import edu.udo.cs.wvtool.util.WVToolLogger;
import edu.udo.cs.wvtool.wordlist.WVTWordList;

public class WVToolsWrapper
{
    WVTContentType   type;
    boolean          parseTag;
    HashSet<String>  idSet;
    WVTWordList      wordList;
    WVTool           wvt;
    WVTConfiguration config;

    private static WVToolsWrapper newInstance(WVTContentType type)
    {
        return new WVToolsWrapper(type);
    }

    public static WVToolsWrapper newInstance()
    {
        return newInstance(WVTContentType.TXT);
    }

    private WVToolsWrapper(WVTContentType type)
    {
        this.type = type;
        switch (type)
        {
            case TXT:
                parseTag = false;
                break;
            case XML:
                parseTag = true;
                break;
            default:
                parseTag = false;
        }
        idSet = new HashSet<String>();
        wvt = new WVTool(true);
        config = getDefaultConfiguration();
        wordList = new WVTWordList(0);
    }

    public WVTConfiguration getDefaultConfiguration()
    {
        if (config == null)
        {
            config = new WVTConfiguration();
            config.setConfigurationRule(WVTConfiguration.STEP_INPUT_FILTER,
                    new WVTConfigurationFact(parseTag ? new XMLInputFilter()
                            : new TextInputFilter()));
            config.setConfigurationRule(WVTConfiguration.STEP_STEMMER,
                    new WVTConfigurationFact(new PorterStemmerWrapper()));
            config.setConfigurationRule(WVTConfiguration.STEP_VECTOR_CREATION,
                    new WVTConfigurationFact(new TFIDF()));
            config.setConfigurationRule(WVTConfiguration.STEP_LOADER,
                    new WVTConfigurationFact(new XMLDocumentLoader()));
        }
        return config;
    }

    public WVTConfiguration getDefaultConfiguration(Writer vsmWriter)
            throws IOException
    {
        WVTConfiguration config = new WVTConfiguration();
        config.setConfigurationRule(WVTConfiguration.STEP_INPUT_FILTER,
                new WVTConfigurationFact(parseTag ? new XMLInputFilter()
                        : new TextInputFilter()));
        config.setConfigurationRule(WVTConfiguration.STEP_STEMMER,
                new WVTConfigurationFact(new PorterStemmerWrapper()));
        config.setConfigurationRule(WVTConfiguration.STEP_VECTOR_CREATION,
                new WVTConfigurationFact(new TFIDF()));
        config.setConfigurationRule(WVTConfiguration.STEP_LOADER,
                new WVTConfigurationFact(new XMLDocumentLoader()));
        WordVectorWriterExtend wvwe = new WordVectorWriterExtend(vsmWriter,
                true);
        config.setConfigurationRule(WVTConfiguration.STEP_OUTPUT,
                new WVTConfigurationFact(wvwe));
        return config;
    }

    public WVTFileInputList createFileList(List<XMLDoc> list)
    {
        WVTFileInputList fileList = new WVTFileInputList(0);
        for (XMLDoc doc : list)
        {
            fileList.addEntry(createWVTXMLDocumentInfo(doc));
        }
        return fileList;
    }

    public WVTXMLDocumentInfo createWVTXMLDocumentInfo(XMLDoc doc)
    {
        return new WVTXMLDocumentInfo(doc.getDocID(), parseTag ? "xml" : "txt",
                "", "english", parseTag ? doc.getRoot().asXML() : doc
                        .getAllText());
    }

    public WVTXMLDocumentInfo createWVTXMLDocumentInfo(String source,
            String context)
    {
        return new WVTXMLDocumentInfo(source, parseTag ? "xml" : "txt", "",
                "english", context);
    }

    public void load(File wordListPath, File idListPath) throws IOException
    {
        Reader reader = new FileReader(wordListPath);
        wordList = new WVTWordList(reader);
        // reader.close();
        reader = new BufferedReader(new FileReader(idListPath));
        String str;
        while ((str = ((BufferedReader) reader).readLine()) != null
                && !str.isEmpty())
        {
            idSet.add(str);
        }
        reader.close();
    }

    public void save(File wordListPath, File idListPath) throws IOException
    {
        Writer writer = new FileWriter(wordListPath);
        wordList.store(writer);
        // writer.flush();
        // writer.close();
        writer = new BufferedWriter(new FileWriter(idListPath));
        for (String str : idSet)
            writer.write(str + "\r\n");
        writer.flush();
        writer.close();
    }

    public void createWVTWordList(List<XMLDoc> list) throws WVToolException
    {
        WVTFileInputList fileList = createFileList(list);
        wordList = wvt.createWordList(fileList, config);
    }

    public WVTWordVector createWordVector(XMLDoc doc, boolean appendWord)
            throws IOException, WVToolException
    {
        String id = doc.getDocID();
        String context = parseTag ? doc.getRoot().asXML() : doc.getAllText();
        return CreateWordVector(id, context, appendWord);
    }

    private WVTWordVector CreateWordVector(String id, String context,
            boolean appendWord) throws IOException, WVToolException
    {
        if (!idSet.contains(id))
            idSet.add(id);
        else
            return createWordVector(context);
        WVTDocumentInfo info = createWVTXMLDocumentInfo(id, context);
        return createWordVector(context, info, config, wordList, appendWord);
    }

    private WVTWordVector createWordVector(String text, WVTDocumentInfo d,
            WVTConfiguration config, WVTWordList wordList, boolean appendWord)
    {

        // Set up the word list properly

        wordList.setAppendWords(appendWord);
        wordList.setUpdateOnlyCurrent(false);

        // Initialize pointers to components for the individual steps

        WVTCharConverter charConverter = null;
        WVTTokenizer tokenizer = null;
        WVTWordFilter wordFilter = null;
        WVTStemmer stemmer = null;
        WVTVectorCreator vectorCreator = null;

        WVTWordVector result = null;

        try
        {

            // Intialize all required components for this document

            charConverter = (WVTCharConverter) config.getComponentForStep(
                    WVTConfiguration.STEP_CHAR_MAPPER, d);
            tokenizer = (WVTTokenizer) config.getComponentForStep(
                    WVTConfiguration.STEP_TOKENIZER, d);
            wordFilter = (WVTWordFilter) config.getComponentForStep(
                    WVTConfiguration.STEP_WORDFILTER, d);
            stemmer = (WVTStemmer) config.getComponentForStep(
                    WVTConfiguration.STEP_STEMMER, d);

            vectorCreator = (WVTVectorCreator) config.getComponentForStep(
                    WVTConfiguration.STEP_VECTOR_CREATION, d);

            // Process the document

            TokenEnumeration tokens = stemmer.stem(wordFilter.filter(tokenizer
                    .tokenize(charConverter.convertChars(
                            new StringReader(text), d), d), d), d);

            while (tokens.hasMoreTokens())
            {
                wordList.addWordOccurance(tokens.nextToken());
            }

            result = vectorCreator.createVector(wordList
                    .getFrequenciesForCurrentDocument(), wordList
                    .getTermCountForCurrentDocument(), wordList, d);

            wordList.closeDocument(d);

        }
        catch (WVToolException e)
        {

            WVToolLogger.getGlobalLogger().logException(
                    "Problems processing document " + d.getSourceName(), e);

            // If errors should not be skip throw an exception
            // if (!skipErrors)
            // throw new WVToolException("Problems processing document " +
            // d.getSourceName(), e);
            // otherwise do nothing and proceed with the next document

        }

        return result;
    }

    public WVTWordVector createWordVector(String context) throws IOException,
            WVToolException
    {
        WVTool wvt = new WVTool(true);
        return wvt.createVector(context, wordList);
    }

    public void pruneByFrequency(int min, int max)
    {
        wordList.pruneByFrequency(min, max);
    }

    public static void getXMLFiles(File f, HashSet<File> set)
    {
        if (f.isFile())
        {
            set.add(f);
            return;
        }
        File[] files = f.listFiles(new FileFilter()
        {
            @Override
            public boolean accept(File dir)
            {
                if (dir.getName().toLowerCase().endsWith(".xml")
                        || dir.isDirectory())
                    return true;
                return false;
            }
        });
        for (File file : files)
            getXMLFiles(file, set);
    }

    public static List<XMLDoc> getXmlList(String root_path)
    {
        final SAXReader reader = new SAXReader();
        // String root_path = getAttribute(DATA, "root_path");
        // String root_path = "D:/INEX Dataset/test set/1484914";
        File root = new File(root_path);
        HashSet<File> set = new HashSet<File>();
        getXMLFiles(root, set);
        System.out.println("total xml:" + set.size());
        List<XMLDoc> list = new ArrayList<XMLDoc>();
        int errorFile = 0;
        for (File file : set)
        {
            try
            {
                Document d = reader.read(file);
                String name = file.getName();
                name = name.substring(0, name.length() - 4);
                // Integer.parseInt(name);
                XMLDoc xmldoc = new XMLDoc(d, name);
                list.add(xmldoc);
            }
            catch (Exception e)
            {
                System.out.println("parse error: " + file.getAbsolutePath());
                errorFile++;
                // e.printStackTrace();
                // System.exit(0);
            }
        }
        System.out.println("error file: " + errorFile);
        return list;
    }

    public static void main(String[] args) throws WVToolException, IOException
    {
        // WVToolsWrapper wrapper =
        // WVToolsWrapper.newInstance(WVTContentType.XML);
        // File wordListPath = new File(getAttribute(SIMILARITY,
        // "word_list_path"));
        // File idListPath = new File(getAttribute(SIMILARITY, "id_list_path"));
        // if (wordListPath.exists() && idListPath.exists())
        // wrapper.load(wordListPath, idListPath);
        // BufferedReader in = new BufferedReader(new
        // InputStreamReader(System.in));
        // // String path = in.readLine();
        // String path = "D:/INEX Dataset/mix set/shakespeare";
        // List<XMLDoc> list = getXmlList(path);
        // int size = list.size();
        // System.out.println(size);
        // for (int i = 0; i < size; i++)
        // {
        // wrapper.createWordVector(list.get(i), true);
        // if (i % 100 == 0)
        // System.out.println(i + "th document parse finished");
        // }
        // System.out.println(wrapper.wordList.getNumDocuments());
        // System.out.println(wrapper.wordList.getNumWords());
        // wrapper.wordList.pruneByFrequency(10, 15000);
        // System.out.println(wrapper.wordList.getNumWords());
        // wrapper.save(wordListPath, idListPath);
    }
}
